#'
#' @title ## 差异基因火山图与热图
#' @TODO  利用异分析结果，并绘制火山图和热图
#' @description 利用异分析结果，并绘制火山图和热图
#' @details 如果`saveplot = T`，则会在路径`output_dir`下生成一个文件夹，以exp的变量名字命名，用于保存结果图片
#' @param DEG_res  一个差异分析结果，表头必须为以下格式。是全部基因的差异分析结果，未经筛选的全部结果。
#'                  列名必须包含 gene、log2FC、pvalue、adjusted_pvalue
#' @param sample_group 在做差异分析时的分组文件，用于热图列注释。需要是tumor与normal，不然会报错。
#' @param exp: 差异分析时的表达谱文件，用于热图列注释，基因行名，样本列名
#' @param topN 热图展示的基因数目，会选取前topN/2个上调与下调组成热图
#' @param saveplot 是否保存图片，默认参数 `F`
#' @param log2FC_value 差异筛选log2FC数值
#' @param adj.p 筛选差异基因 矫正后p值阈值
#' @param output_dir 生成图片文件目录
#' @param display_gene 火山图中展示的基因，是一个字符串向量,会在热图中展示显著的基因
#' @param var_name 在生成文件夹以及命名的时候用于命名
#' @param cluster_name 热图组别标签命名
#'
#' @return *list*
#' @return  一个list，三个图，火山图，热图，合图，卡阈值后的差异分析结果
#'
#' @examples d <- volcano_heatmap(DEG_res = degs_all, exp = brca_exp, sample_group = deg_group, log2FC_value = 0.585, saveplot = T, topN = 40)
#'
#' @export
#' @author *WYK*
#'
DEG_volcano_heatmap <- function(DEG_res = NULL, exp = NULL, sample_group = NULL,
                                log2FC_value = 1, saveplot = F, topN = 50, adj.p = 0.05, p = NULL,
                                output_dir = "/pub/users/innertech/wyk/Project_output/GAP1234/",
                                display_gene = NULL, var_name = NULL, color_used = NULL, cluster_name = "group", color_in_vol = NULL) {
  if (length(var_name) == 0) {
    var_name <- deparse(substitute(exp))
  }

  figure_list <- list()

  DEG_res <- as.data.frame(DEG_res)
  exp <- as.data.frame(exp)
  sample_group <- as.data.frame(sample_group)

  require(tidyverse)
  require(cowplot)

  if (!is.null(p)) {
    DEG_res <- DEG_res %>%
      select(-adjusted_pvalue) %>%
      dplyr::rename(adjusted_pvalue = pvalue)
    adj.p <- p
  }

  degs <- DEG_res %>%
    filter(adjusted_pvalue < adj.p) %>%
    filter(log2FC > log2FC_value | log2FC < (-log2FC_value))

  Down <- degs %>%
    filter(log2FC < (-log2FC_value)) %>%
    nrow(.)

  Up <- degs %>%
    filter(log2FC > log2FC_value) %>%
    nrow(.)

  NoSig <- nrow(DEG_res) - Up - Down

  up_chr <- sprintf("Up: %s", Up)
  down_chr <- sprintf("Down: %s", Down)

  threshold <- DEG_res %>%
    mutate(threshold = case_when(
      log2FC > log2FC_value & adjusted_pvalue < adj.p ~ 'Up', # ,sprintf("Up: %s", Up)
      log2FC < -log2FC_value & adjusted_pvalue < adj.p ~ 'Down', # ,sprintf("Down: %s", Down)
      adjusted_pvalue > adj.p ~ "NoSig", # sprintf("NoSig: %s", NoSig),
      between(log2FC, -log2FC_value, log2FC_value) ~ "NoSig" # sprintf("NoSig: %s", NoSig)
    )) %>%
    select(threshold) %>%
    pull() %>%
    factor(.)

  Volcano_plot <- ggplot(data = DEG_res, aes(x = log2FC, y = -log10(adjusted_pvalue))) +
    xlab("log<sub>2</sub>Fold Change") +
    ylab("-log<sub>10</sub>adj.pvalue") +
    geom_point(size = 1, alpha = .9, aes(colour = threshold)) +
    theme_test(base_size = 15) +
    scale_color_manual(values = c(
      'Down' = "#367EB8",
      "NoSig" = "grey",
      'Up' = "#E4191C"
    )) +
    geom_hline(yintercept = -log10(adj.p), lty = 2, colour = "#242424") +
    theme(
      legend.position = "top",
      legend.title = element_blank(),
      legend.text = element_text(size = 10)
    ) +
    theme(
      axis.title.x = ggtext::element_markdown(),
      axis.title.y = ggtext::element_markdown()
    )

  if (!is.null(p)) {
    Volcano_plot <- Volcano_plot + ylab("-log<sub>10</sub>pvalue")
  }

  # if (as.numeric(log2FC_value) == 0) {
  #   Volcano_plot
  # } else {
  Volcano_plot <- Volcano_plot +
    geom_vline(xintercept = -log2FC_value, lty = 2, colour = "#242424") +
    geom_vline(xintercept = log2FC_value, lty = 2, colour = "#242424")
  # }

  Down_gene <- DEG_res %>%
    filter(adjusted_pvalue < adj.p) %>%
    filter(log2FC > log2FC_value | log2FC < -log2FC_value) %>%
    arrange(log2FC) %>%
    select(gene) %>%
    pull() %>%
    .[1:as.integer(topN / 2)]

  Top_gene <- DEG_res %>%
    filter(log2FC > log2FC_value | log2FC < -log2FC_value) %>%
    filter(adjusted_pvalue < adj.p) %>%
    arrange(desc(log2FC)) %>%
    select(gene) %>%
    pull() %>%
    .[1:as.integer(topN / 2)]

  topgene <- c(Top_gene, Down_gene)

  forheatmap_exp <- exp[topgene, ]

  if (!is.null(display_gene)) {
    forheatmap_exp <- na.omit(exp[display_gene, ])

    display_gene <- display_gene[display_gene %in% {
      degs %>% pull(1)
    }]

    Volcano_plot <- Volcano_plot + ggrepel::geom_text_repel(
      data = DEG_res %>% filter(gene %in% display_gene),
      aes(label = gene), max.overlaps = 10
    ) + geom_point(
      data = DEG_res %>% filter(gene %in% display_gene),
      color = "black", shape = 21
    )
  }

  p_heatmap <- tryCatch(
    Heatmap_manul_NoSplit_forVol(
      data_input = forheatmap_exp,
      color_used = color_used,
      cluster_name = cluster_name,
      group_infor = sample_group %>% arrange(group) %>% column_to_rownames("sample"),
      saveplot = F,
      output_dir = "./",
      var_name = var_name,
      width_used = 4.7, height_used = 6, heatmap_name = " "
    ),
    error = function(e) {
      return("error in heatmap")
    }
  )

  if (class(p_heatmap)[1] == "character") {
    p_heatmap <- ggplot()
  }

  p <- plot_grid(Volcano_plot, p_heatmap, rel_widths = c(1.5, 2), labels = "AUTO", label_size = 20, scale = c(1, .98))

  # print(p)

  figure_list[[1]] <- Volcano_plot
  figure_list[[2]] <- p_heatmap
  figure_list[[3]] <- p

  message(sprintf("There're %s Differential Genes.", nrow(degs)), sprintf("%s is Down.", Down), sprintf("%s is Up.", Up))

  figure_list[[4]] <- degs

  if (saveplot) {
    if (!dir.exists(sprintf("%s/DEGS_Vol_heatmap_%s", output_dir, var_name))) {
      dir.create(sprintf("%s/DEGS_Vol_heatmap_%s", output_dir, var_name), showWarnings = F, recursive = T)
    } else {
      print(sprintf("Dir '%s/DEGS_Vol_heatmap_%s' is existed.", output_dir, var_name))
    }

    ggsave2(
      filename = sprintf("%s/DEGS_Vol_heatmap_%s/Volcano_plot_%s.pdf", output_dir, var_name, var_name),
      plot = Volcano_plot, width = 4.43, height = 5
    )
    ggsave2(
      filename = sprintf("%s/DEGS_Vol_heatmap_%s/p_heatmap_%s.pdf", output_dir, var_name, var_name),
      plot = p_heatmap, width = 4.43, height = 5
    )

    ggsave2(
      filename = sprintf("%s/DEGS_Vol_heatmap_%s/Vol_heatmap_%s.pdf", output_dir, var_name, var_name),
      plot = p, width = 8.7, height = 5
    )
  }

  return(figure_list)
}

Heatmap_manul_NoSplit_forVol <- function(data_input = NULL, color_used = NULL, cluster_name = "cluster", group_infor = NULL,
                                         saveplot = T, output_dir = "./", var_name = NULL, width_used = 9, height_used = 9, heatmap_name = " ") {
  if (is.null(var_name)) {
    var_name <- paste0("_", paste0(sample(letters, 4), collapse = "", sep = ""))
  }
  group_infor_2 <- na.omit(group_infor)
  group_infor <- na.omit(group_infor)

  library(tidyverse)
  library(ComplexHeatmap)
  library(cowplot)
  library(ggpubr)

  sample_common <- intersect(rownames(group_infor), colnames(data_input))
  group_infor <- group_infor[sample_common, , drop = F]
  data_input <- data_input[, sample_common]

  group_chara <- sort(unique(as.character(group_infor$group)))

  if (is.null(color_used)) {
    col_name <- c("#7FC97F", "#BEAED4") #  RColorBrewer::brewer.pal(8, "Set1")[1:length(group_chara)]
    col_name <- col_name[seq_along(unique(group_infor %>% pull(group)))]
  } else {
    col_name <- color_used[seq_along(unique(group_infor %>% pull(group)))]
  }

  names(col_name) <- group_chara

  if (ncol(group_infor) >= 2) {
    col_anno <- HeatmapAnnotation(
      df = group_infor, # %>% .[, "group", drop = F]
      col = list(
        group = col_name
      ),
      annotation_legend_param = list(group = list(title = c(cluster_name, colnames(group_infor)[2:ncol(group_infor)]))),
      annotation_name_side = "left",
      annotation_label = c(cluster_name, colnames(group_infor)[2:ncol(group_infor)])
    )
  } else {
    col_anno <- HeatmapAnnotation(
      df = group_infor, # %>% .[, "group", drop = F]
      col = list(
        group = col_name
      ),
      annotation_name_side = "right",
      annotation_label = cluster_name,
      annotation_legend_param = list(group = list(title = cluster_name))
    )
  }

  data_input_scaled <- t(apply(data_input, 1, function(x) scale(x, center = T, scale = T)))
  colnames(data_input_scaled) <- colnames(data_input)
  all(colnames(data_input_scaled) == group_infor %>% rownames())

  group_infor_2 <- group_infor_2[rownames(group_infor_2) %in% rownames(group_infor), , drop = F]

  col_zscore <- circlize::colorRamp2(c(-2, 0, 2), c("navy", "white", "firebrick3")) # c("#0a5aa5", "white", "firebrick3")
  # col_zscore <- circlize::colorRamp2(c(-2, 0, 2), c("purple", "black", "yellow"))

  cluster_res <- Heatmap(
    matrix = data_input_scaled,
    name = heatmap_name,
    col = col_zscore,
    show_column_dend = F,
    show_column_names = F,
    show_row_names = T,
    cluster_rows = T,
    show_row_dend = F,
    # rect_gp = gpar(col = "white", lwd = 1),
    row_names_gp = gpar(fontsize = 6),
    column_title = " ",
    top_annotation = col_anno,
    column_gap = unit(0, "mm"),
    # column_order = rownames(group_infor_2),
    # heatmap_legend_param = list(direction = "horizontal"),
    column_split = group_infor %>% .[, 1, drop = F],
    border = F
    # left_annotation = row_anno,
    # row_names_side = "right"
  )

  gg.cluster_res <- cluster_res %>%
    draw(.,
      merge_legend = T, heatmap_legend_side = "right",
      annotation_legend_side = "right"
    ) %>%
    grid.grabExpr() %>%
    ggplotify::as.ggplot() +
    theme(
      plot.background = element_rect(fill = "white", color = "white"),
      plot.margin = unit(c(.1, .1, .3, .1), "cm")
    )

  if (saveplot) {
    dir_now <- str_glue("{output_dir}/output/heatmap/")

    if (!dir.exists(dir_now)) {
      dir.create(dir_now, recursive = T)
    } else {
      print(str_c(dir_now, " is ready."))
    }

    ggsave2(
      filename = str_glue("{dir_now}/figure_heat_plot{var_name}.pdf"),
      plot = gg.cluster_res, width = width_used, height = height_used
    )
    # ggsave2(
    #     filename = str_glue("{dir_now}/figure_heat_plot{var_name}.tiff"),
    #     plot = gg.cluster_res, width = width_used, height = height_used, dpi = 300
    # )
    # ggsave2(
    #     filename = str_glue("{dir_now}/figure_heat_plot{var_name}_dpi72.tiff"),
    #     plot = gg.cluster_res, width = width_used, height = height_used, dpi = 72
    # )
  }
  return(gg.cluster_res)
}
